from collections import Counter
import math


def cosine_ngram_similarity(s1, s2, n=2, is_char=True):
    # 生成N-gram并计算词频
    def generate_ngrams(s, n, is_char):
        if is_char:
            return [s[i:i + n].lower() for i in range(len(s) - n + 1)]
        else:
            return s.lower().split()

    ngrams1 = generate_ngrams(s1, n, is_char)
    ngrams2 = generate_ngrams(s2, n, is_char)

    # 计算词频向量
    vector1 = Counter(ngrams1)
    vector2 = Counter(ngrams2)

    # 获取所有N-gram的并集
    all_ngrams = set(vector1.keys()).union(set(vector2.keys()))

    # 计算余弦相似度
    dot_product = sum(vector1.get(ngram, 0) * vector2.get(ngram, 0) for ngram in all_ngrams)
    norm1 = math.sqrt(sum(v ** 2 for v in vector1.values()))
    norm2 = math.sqrt(sum(v ** 2 for v in vector2.values()))

    return dot_product / (norm1 * norm2) if (norm1 * norm2) != 0 else 0


s1 = "主键ID"
s2 = "ID"
similarity = cosine_ngram_similarity(s1, s2, n=2, is_char=True)
print(f"字符级N-gram余弦相似度 (n=2): {similarity:.2f}")
